# encoding=utf-8
import threading

import cv2
import torch
from networks.whole_network import NF as Mymodel
from PIL import Image
from torchvision import transforms
import numpy as np


def array2bytes(array_img, suffix):
    # 对数组的图片格式进行编码
    success, encoded_array = cv2.imencode("." + suffix, array_img)
    # 将数组转为bytes
    bytes_img = encoded_array.tostring()

    return bytes_img


class CeilNerveformerWrapper:
    _inst = None
    _lock = threading.Lock()

    def __init__(self):
        self.model = Mymodel()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        model_d = torch.load("ceill-nerveformer.pth")
        self.model.load_state_dict(model_d)
        self.model.to(self.device)
        self.model.eval()
        self.transform = transforms.Compose([
            transforms.Resize([384, 384]),
            transforms.ToTensor()
        ])

    def _preprocess(self, data_in: str):
        image = Image.open(data_in)
        ori_size = image.size
        gray = image.convert('L')
        trans = self.transform(gray)
        trans = trans.unsqueeze(0)
        return trans, ori_size

    @torch.no_grad()
    def infer(self, image_path: str):
        trans_dim, ori_size = self._preprocess(image_path)
        trans_cuda = trans_dim.to(self.device)
        res = self.infer_wrapper(trans_cuda, ori_size)
        return res

    def infer_wrapper(self, cuda_tensor, ori_size):
        model_ceil, model_nerve = self.model(cuda_tensor)
        ceil_mask = self.make_mask(model_ceil, ori_size)
        nerve_mask = self.make_mask(model_nerve, ori_size)
        return ceil_mask, nerve_mask

    def make_mask(self, model_out, ori_size):
        cpu_out = model_out.detach().cpu().numpy()
        mask = self.return_prediction(cpu_out)
        W, H = ori_size
        img_arr = cv2.resize(mask, (W, H), interpolation=cv2.INTER_CUBIC)
        img_bytes = array2bytes(img_arr, "png")
        image = np.asarray(bytearray(img_bytes), dtype=np.uint8)
        img_arr = cv2.imdecode(image, cv2.IMREAD_COLOR)
        return img_arr

    def return_prediction(self, pred):
        mask = pred.squeeze(0).squeeze(0)
        mask = (mask * 255).astype(np.uint8)
        return mask

    def __new__(cls, *args, **kwargs):
        CeilNerveformerWrapper._lock.acquire()
        if cls._inst is None:
            cls._inst = super(CeilNerveformerWrapper, cls).__new__(cls)
        CeilNerveformerWrapper._lock.release()
        return cls._inst